import asyncio
import edge_tts
from pygame import mixer
import io

# 配置参数
TEXT = "你好，这是一个使用edge-tts进行语音合成并实时播放的示例。"
VOICE = "zh-CN-XiaoxiaoNeural"  # 中文女声
OUTPUT_FILE = "output.mp3"  # 可选：保存音频到文件


async def main():
    # 初始化pygame混音器
    mixer.init()

    # 创建TTS对象
    communicate = edge_tts.Communicate(TEXT, VOICE)

    # 异步获取音频块并播放
    audio_buffer = io.BytesIO()
    async for chunk in communicate.stream():
        if chunk["type"] == "audio":
            # 将音频块写入缓冲区
            audio_buffer.write(chunk["data"])

            # 当缓冲区达到一定大小时开始播放
            if audio_buffer.tell() > 1024 * 10:  # 10KB
                # 保存当前位置
                current_pos = audio_buffer.tell()

                # 回退到开始并加载到混音器
                audio_buffer.seek(0)
                mixer.music.load(audio_buffer)
                mixer.music.play()
                # 等待播放完成
                while mixer.music.get_busy():
                    await asyncio.sleep(0.1)
                # 重置缓冲区
                audio_buffer = io.BytesIO()
                # 重置缓冲区

    # 播放剩余的音频
    if audio_buffer.tell() > 0:
        audio_buffer.seek(0)
        mixer.music.load(audio_buffer)
        mixer.music.play()

    # 等待播放完成
    while mixer.music.get_busy():
        await asyncio.sleep(0.1)

    # 保存音频到文件（可选）
    if OUTPUT_FILE:
        with open(OUTPUT_FILE, "wb") as f:
            audio_buffer.seek(0)
            f.write(audio_buffer.read())

    # 清理资源
    mixer.quit()


if __name__ == "__main__":
    asyncio.run(main())